#include "DynamicLinkedLibraryInjector.h"

DWORD NahidaProject::DynamicLibraryInjector::GetProcessID(const std::string& processName){
	PROCESSENTRY32 proc_entry;
	proc_entry.dwSize = sizeof(proc_entry);

	auto snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, NULL);
	if (!snapshot) { 
		return GetLastError(); 
	}

	Process32First(snapshot, &proc_entry);
	if (!processName.compare(proc_entry.szExeFile)) {
		CloseHandle(snapshot);
		return proc_entry.th32ProcessID;
	}

	while (Process32Next(snapshot, &proc_entry)) {
		if (!processName.compare(proc_entry.szExeFile)) {
			CloseHandle(snapshot);
			return proc_entry.th32ProcessID;
		}
	}

	CloseHandle(snapshot);
	return 0;
}

bool NahidaProject::DynamicLibraryInjector::Inject(const std::string& processName, const std::string& dynamicLinkedLibraryName){

	auto pid = GetProcessID(processName);
	if (!pid) { 
		return false;
	}

	auto proc = OpenProcess(PROCESS_ALL_ACCESS, NULL, pid);
	if (!proc) { 
		return false;
	}

	auto kernel32 = GetModuleHandleA("kernel32.dll");
	if (!kernel32) { 
		return false; 
	}

	auto load_library = GetProcAddress(kernel32, "LoadLibraryW");
	if (!load_library) { 
		return false;
	}

	char* dll_path = nullptr;
	GetFullPathName(dynamicLinkedLibraryName.c_str(), MAX_PATH, dll_path, NULL);

	auto alloc = VirtualAllocEx(proc, NULL, sizeof(dll_path), MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
	if (!alloc) { 
		return false;
	}

	WriteProcessMemory(proc, alloc, dll_path, sizeof(dll_path), NULL);
	auto thread = CreateRemoteThread(proc, NULL, NULL, (LPTHREAD_START_ROUTINE)load_library, alloc, NULL, NULL);
	if (!thread) { 
		return false;
	}

	WaitForSingleObject(thread, INFINITE);
	CloseHandle(proc);
	return true;
}
